package org.example.tool;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.nio.charset.Charset;
import java.util.*;

public class HuffmanTree {
    public static final Charset CHARSET = Charset.defaultCharset();
    private Node root;
    public Node getRoot(){
        return root;
    }
    public HuffmanTree(){}
    public HuffmanTree(Node root) {
        this.root = root;
    }
    public static class Node {
        protected Node one, zero;
        protected char c;
        private boolean isLeaf;
        Node(){}
        Node(char c) {
            this.isLeaf = true;
            this.c = c;
        }
        public Node getOne() {
            return one;
        }
        public Node getZero() {
            return zero;
        }
        public char getChar() {
            return c;
        }
        public boolean isLeaf() {
            return isLeaf;
        }
    }
    private static class NodeSortedList {
        static class WeightedNode extends Node {
            int weight;
            WeightedNode(int weight, Node one, Node zero){
                super();
                this.weight = weight;
                this.one = one;
                this.zero = zero;
            }
            WeightedNode(char c, int weight) {
                super(c);
                this.weight = weight;
            }
            static WeightedNode union(WeightedNode a, WeightedNode b){
                return new WeightedNode(a.weight + b.weight, a, b);
            }
        }
        private List<WeightedNode> all = new ArrayList<>();
        public void add(char c, int weight) {
            add(new WeightedNode(c, weight));
        }
        private void add(WeightedNode n) {
//            int index = Collections.binarySearch(all, n, Comparator.comparingInt(a -> a.weight));
            int l = 0, r = all.size(), mid = (l + r) / 2;
            while(l<r){
                if(all.get(mid).weight < n.weight) {
                    r = mid;
                } else if(all.get(mid).weight > n.weight) {
                    l = mid + 1;
                } else {
                    break;
                }
                mid = (l + r) / 2;
            }
            all.add(mid, n);
        }
        public Node reduceAll(){
            while (all.size()>1) {
                WeightedNode union = WeightedNode.union(all.get(all.size() - 1), all.get(all.size() - 2));
                this.all = this.all.subList(0, all.size()-2);
                this.add(union);
            }
            return all.get(0);
        }
    }
    public static class FromTextBuilder {
        private Map<Character, Integer> map;
        public FromTextBuilder() {
            this.map = new HashMap<>();
        }
        public void addString(String str) {
            for(int i = 0; i<str.length(); i++){
                addChar(str.charAt(i));
            }
        }
        public void addChar(char c) {
            int count = this.map.getOrDefault(c, 0);
            this.map.put(c, count+1);
        }
        public HuffmanTree build() {
            NodeSortedList all = new NodeSortedList();
            for(Character c: map.keySet()) {
                all.add(c, map.get(c));
            }
            Node node = all.reduceAll();
            HuffmanTree huffmanTree = new HuffmanTree();
            huffmanTree.root = node;
            return huffmanTree;
        }
    }

    public String toString(){
        var b = new ByteArrayOutputStream();
        HuffmanTreeOutputStream huffmanTreeWriter = new HuffmanTreeOutputStream(b);
        try {
            huffmanTreeWriter.writeTree(this);
            huffmanTreeWriter.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return b.toString();
    }

}
